// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

//! A safe, idiomatic Rust interface for dynamic library loading.
//!
//! On Unix-like systems, this module wraps `libdl`. It is designed to be safe
//! against `LD_PRELOAD`-based interception by resolving the real `dl*` functions
//! at runtime.
#![cfg(unix)]

use std::{
    ffi::{CStr, CString, OsStr, OsString},
    mem,
    os::unix::ffi::OsStrExt,
    path::{Path, PathBuf},
    ptr,
    sync::Arc,
};

use bitflags::bitflags;
use libc::{c_char, c_int, c_void};
use once_cell::sync::Lazy;
use thiserror::Error;

/// A secure FFI layer for libdl that bypasses `LD_PRELOAD` hooks.
///
/// This module resolves the true `dl*` functions at runtime, avoiding local
/// function interception.
pub mod ffi {
    use super::*;

    // Re-export constants from libc for use in flags.
    pub use libc::{Dl_info, RTLD_GLOBAL, RTLD_LAZY, RTLD_LOCAL, RTLD_NEXT, RTLD_NOW};

    #[cfg(target_os = "linux")]
    pub use libc::{RTLD_DEEPBIND, RTLD_NODELETE, RTLD_NOLOAD};

    // Type aliases for the dl* function pointers.
    type DlOpenFn = unsafe extern "C" fn(*const c_char, c_int) -> *mut c_void;
    type DlCloseFn = unsafe extern "C" fn(*mut c_void) -> c_int;
    type DlSymFn = unsafe extern "C" fn(*mut c_void, *const c_char) -> *mut c_void;
    type DlAddrFn = unsafe extern "C" fn(*const c_void, *mut Dl_info) -> c_int;
    type DlErrorFn = unsafe extern "C" fn() -> *mut c_char;

    // Use `once_cell::Lazy` to resolve the real function pointers only once.
    static DLOPEN_FN: Lazy<DlOpenFn> = Lazy::new(|| bootstrap_symbol("dlopen"));
    static DLERROR_FN: Lazy<DlErrorFn> = Lazy::new(|| bootstrap_symbol("dlerror"));
    static DLSYM_FN: Lazy<DlSymFn> = Lazy::new(|| bootstrap_symbol("dlsym"));
    static DLADDR_FN: Lazy<DlAddrFn> = Lazy::new(|| bootstrap_symbol("dladdr"));
    static DLCLOSE_FN: Lazy<DlCloseFn> = Lazy::new(|| bootstrap_symbol("dlclose"));

    /// Fetches a function pointer from the next library in the link order.
    /// This bypasses `LD_PRELOAD` hooks by searching after the current module.
    fn bootstrap_symbol<T>(name: &str) -> T {
        let symbol = CString::new(name).expect("FATAL: bootstrap symbol name contains null byte.");

        // SAFETY: `libc::dlsym` is called with `RTLD_NEXT`, a valid pseudo-handle,
        // and a valid C-string. This is a standard and safe way to look up symbols.
        let ptr = unsafe { libc::dlsym(RTLD_NEXT, symbol.as_ptr()) };
        if ptr.is_null() {
            panic!("FATAL: Failed to resolve bootstrap function '{}'.", name);
        }
        // SAFETY: The caller of `bootstrap_symbol` (the `Lazy` initializers) guarantees
        // that `T` is the correct function signature for the symbol name.
        unsafe { mem::transmute_copy(&ptr) }
    }

    /// Opens a dynamic library.
    ///
    /// # Safety
    /// `filename` must be a valid, null-terminated C string. `flag` must be valid.
    #[inline]
    pub unsafe fn dlopen(filename: *const c_char, flag: c_int) -> *mut c_void {
        // SAFETY: Caller guarantees valid arguments. `DLOPEN_FN` is a valid pointer resolved at init.
        unsafe { DLOPEN_FN(filename, flag) }
    }

    /// Returns the last error from a dynamic linking operation.
    ///
    /// # Safety
    /// The C `dlerror` is not thread-safe. This wrapper is also not guaranteed to be
    /// thread-safe if called concurrently with other dl* functions in the same thread.
    #[inline]
    pub unsafe fn dlerror() -> *mut c_char {
        // SAFETY: Caller ensures no data races. `DLERROR_FN` is a valid pointer.
        unsafe { DLERROR_FN() }
    }

    /// Finds a symbol in a dynamic library.
    ///
    /// # Safety
    /// `handle` must be valid. `symbol` must be a valid, null-terminated C string.
    #[inline]
    pub unsafe fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void {
        // SAFETY: Caller guarantees valid arguments. `DLSYM_FN` is a valid pointer.
        unsafe { DLSYM_FN(handle, symbol) }
    }

    /// Translates an address to symbolic information.
    ///
    /// # Safety
    /// `addr` can be any pointer. `info` must point to a valid, mutable `Dl_info` struct.
    #[inline]
    pub unsafe fn dladdr(addr: *const c_void, info: *mut Dl_info) -> c_int {
        // SAFETY: Caller guarantees valid arguments. `DLADDR_FN` is a valid pointer.
        unsafe { DLADDR_FN(addr, info) }
    }

    /// Closes a dynamic library.
    ///
    /// # Safety
    /// `handle` must be a valid pointer from a successful `dlopen` call, or `RTLD_NEXT` etc.
    #[inline]
    pub unsafe fn dlclose(handle: *mut c_void) -> c_int {
        // SAFETY: Caller guarantees a valid handle. `DLCLOSE_FN` is a valid pointer.
        unsafe { DLCLOSE_FN(handle) }
    }
}

/// An error type for dynamic linking operations.
#[derive(Debug, Error)]
pub enum DlError {
    /// The provided library path contains an invalid internal null byte.
    #[error("Path '{}' contains a null byte", path.display())]
    InvalidPath { path: PathBuf },

    /// The provided symbol name contains an invalid internal null byte.
    #[error("Symbol '{}' contains a null byte", name.to_string_lossy())]
    InvalidSymbolName { name: OsString },

    /// The given address could not be mapped to any known library by `dladdr`.
    #[error("Address {addr:p} is not associated with any known library")]
    InvalidAddress { addr: *const c_void },

    /// Failed to load the dynamic library due to a specific system error.
    #[error("Failed to load library '{}': {}", path.display(), reason.to_string_lossy())]
    LoadFailed { path: PathBuf, reason: OsString },

    /// `dladdr` failed to look up an address due to a specific system error.
    #[error("Failed to look up address {addr:p}: {}", reason.to_string_lossy())]
    AddressLookupError {
        addr: *const c_void,
        reason: OsString,
    },

    /// `dlsym` failed to look up a symbol due to a specific system error.
    #[error("Failed to lookup symbol '{}': {}", name.to_string_lossy(), reason.to_string_lossy())]
    SymbolLookupError { name: OsString, reason: OsString },
}

/// An inner RAII guard that owns the raw library handle and calls `dlclose` on drop.
#[derive(Debug)]
#[repr(transparent)]
struct DlGuard(ptr::NonNull<c_void>);

impl DlGuard {
    /// Returns the raw, untyped handle to the library.
    #[inline]
    pub fn as_raw(&self) -> *mut c_void {
        self.0.as_ptr()
    }
}

impl Drop for DlGuard {
    fn drop(&mut self) {
        let raw_handle = self.as_raw();
        // Do not close the special `RTLD_NEXT` pseudo-handle.
        if raw_handle != ffi::RTLD_NEXT {
            // SAFETY: `raw_handle` is valid because it comes from a successful `dlopen` call,
            // and this `drop` is called only once when the `Arc`'s ref-count drops to zero.
            unsafe { ffi::dlclose(raw_handle) };
        }
    }
}

// SAFETY: A library handle is safe to send between threads.
unsafe impl Send for DlGuard {}
// SAFETY: A library handle is safe to access from multiple threads.
unsafe impl Sync for DlGuard {}

/// A reference-counted, RAII handle to a loaded dynamic library.
///
/// This handle is cheap to clone and manages the library's lifetime via `Arc`.
/// The library is automatically closed only when the last handle is dropped.
#[derive(Debug, Clone)]
pub struct DlHandle {
    guard: Arc<DlGuard>,
}

impl DlHandle {
    /// Finds a symbol by name within this library. A convenience wrapper around `dlsym`.
    #[inline]
    pub fn find_symbol<S: AsRef<OsStr>>(&self, name: S) -> Result<Option<DlSymbol>, DlError> {
        dlsym(self, name)
    }

    /// Returns the raw, untyped handle to the library.
    #[inline]
    pub fn as_raw(&self) -> *mut c_void {
        self.guard.as_raw()
    }
}

/// A special pseudo-handle for `dlsym` to search in libraries loaded after this one.
pub static NEXT_HANDLE: Lazy<DlHandle> = Lazy::new(|| DlHandle {
    guard: Arc::new(DlGuard(unsafe {
        ptr::NonNull::new_unchecked(ffi::RTLD_NEXT)
    })),
});

/// A handle to a symbol within a dynamic library.
///
/// This struct is self-contained. It holds a reference-counted handle (`DlHandle`)
/// to the library, ensuring the symbol pointer remains valid for its entire lifetime.
#[derive(Debug, Clone)]
pub struct DlSymbol {
    _handle: DlHandle,
    ptr: ptr::NonNull<c_void>,
}

impl DlSymbol {
    /// Casts the raw symbol pointer to a specific type `T` (e.g., a function pointer).
    ///
    /// # Safety
    /// The caller must ensure `T` is the correct type for the symbol. Casting to a
    /// mismatched type will cause undefined behavior.
    #[inline]
    pub unsafe fn cast<T>(&self) -> T {
        // SAFETY: The caller upholds the contract. This is a standard way to convert a data pointer.
        unsafe { mem::transmute_copy(&self.as_raw()) }
    }

    /// Returns the raw, untyped pointer to the symbol.
    #[inline]
    pub fn as_raw(&self) -> *mut c_void {
        self.ptr.as_ptr()
    }
}

// SAFETY: A symbol pointer is safe to send between threads.
unsafe impl Send for DlSymbol {}
// SAFETY: A symbol pointer can be safely accessed from multiple threads.
unsafe impl Sync for DlSymbol {}

/// A safe wrapper around `ffi::Dl_info` returned by `dladdr`.
#[derive(Debug, Clone, Copy)]
#[repr(transparent)]
pub struct DlInfo(ffi::Dl_info);

impl DlInfo {
    /// Returns the path of the shared object containing the address.
    #[inline]
    pub fn file_name(&self) -> Option<&CStr> {
        if self.0.dli_fname.is_null() {
            return None;
        }
        // SAFETY: `dli_fname` is valid if non-null, as per `dladdr` documentation.
        Some(unsafe { CStr::from_ptr(self.0.dli_fname) })
    }

    /// Returns the base address of the shared object.
    #[inline]
    pub fn base_addr(&self) -> Option<ptr::NonNull<c_void>> {
        ptr::NonNull::new(self.0.dli_fbase)
    }

    /// Returns the name of the nearest symbol to the address.
    #[inline]
    pub fn symbol_name(&self) -> Option<&CStr> {
        if self.0.dli_sname.is_null() {
            return None;
        }
        // SAFETY: `dli_sname` is valid if non-null, as per `dladdr` documentation.
        Some(unsafe { CStr::from_ptr(self.0.dli_sname) })
    }

    /// Returns the address of the nearest symbol.
    #[inline]
    pub fn symbol_addr(&self) -> Option<ptr::NonNull<c_void>> {
        ptr::NonNull::new(self.0.dli_saddr)
    }
}

impl Default for DlInfo {
    fn default() -> Self {
        // Create a zeroed/nulled struct, which is the expected initial state for `dladdr`.
        let dl_info = ffi::Dl_info {
            dli_fname: ptr::null(),
            dli_fbase: ptr::null_mut(),
            dli_sname: ptr::null(),
            dli_saddr: ptr::null_mut(),
        };
        Self(dl_info)
    }
}

bitflags! {
    /// Flags for `dlopen` that control symbol resolution and visibility.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct DlOpenFlag: c_int {
        /// Perform lazy binding. (Default)
        const LAZY = ffi::RTLD_LAZY;
        /// Resolve all symbols at load time.
        const NOW = ffi::RTLD_NOW;
        /// Make symbols from this library globally available.
        const GLOBAL = ffi::RTLD_GLOBAL;
        /// Symbols are not made available to other libraries. (Default)
        const LOCAL = ffi::RTLD_LOCAL;

        /// Prioritize this library's symbols during lookup. (GNU extension)
        #[cfg(target_os = "linux")]
        const DEEP_BIND = ffi::RTLD_DEEPBIND;
        /// Do not unload the library when `dlclose` is called. (GNU extension)
        #[cfg(target_os = "linux")]
        const NO_DELETE = ffi::RTLD_NODELETE;
        /// Test if a library is already loaded without actually loading it. (GNU extension)
        #[cfg(target_os = "linux")]
        const NO_LOAD = ffi::RTLD_NOLOAD;
    }
}

impl Default for DlOpenFlag {
    /// The default flags for `dlopen` are `LAZY`.
    fn default() -> Self {
        DlOpenFlag::LAZY
    }
}

/// Opens a dynamic library with the specified flags.
///
/// # Errors
/// Returns `DlError::InvalidPath` or `DlError::LoadFailed` on failure.
#[inline]
pub fn dlopen<P>(path: P, flags: DlOpenFlag) -> Result<DlHandle, DlError>
where
    P: AsRef<Path>,
{
    let path_ref = path.as_ref();
    let path_bytes = path_ref.as_os_str().as_bytes();
    let filename = CString::new(path_bytes).map_err(|_| DlError::InvalidPath {
        path: path_ref.to_owned(),
    })?;

    // Clear any existing error state.
    let _ = dlerror();

    // SAFETY: `filename` is a valid C-string, and `flags` are valid bitflags.
    // The call is to our securely resolved `dlopen` function.
    let handle = unsafe { ffi::dlopen(filename.as_ptr(), flags.bits()) };

    match ptr::NonNull::new(handle) {
        Some(handle) => Ok(DlHandle {
            guard: Arc::new(DlGuard(handle)),
        }),
        None => Err(DlError::LoadFailed {
            path: path_ref.to_owned(),
            reason: dlerror().unwrap_or_else(|| OsStr::new("unknown error").to_owned()),
        }),
    }
}

/// Returns the last error message from a dynamic linking operation as a string.
/// This should be called immediately after a `dl*` function fails.
#[inline]
pub fn dlerror() -> Option<OsString> {
    // SAFETY: `ffi::dlerror` is safe to call.
    let ptr = unsafe { ffi::dlerror() };
    if ptr.is_null() {
        return None;
    }

    // SAFETY: A non-null return from `dlerror` points to a valid, static C-string.
    // We must copy it into an owned `OsString` before the next `dl*` call overwrites it.
    let byte_slice = unsafe { CStr::from_ptr(ptr) }.to_bytes();
    let err_msg = OsStr::from_bytes(byte_slice);

    Some(err_msg.to_owned())
}

/// Finds the address of a symbol within a given library handle.
///
/// This function distinguishes between a symbol that is not found and an error
/// that occurs during the lookup process, as per the `dlsym` specification.
///
/// # Returns
///
/// *   `Ok(Some(DlSymbol))` if the symbol is found successfully. The returned `DlSymbol`
///     ensures the library remains loaded for as long as the symbol is in scope.
/// *   `Ok(None)` if the symbol is not found in the specified library. This is not
///     considered an error.
/// *   `Err(DlError::SymbolLookupError)` if an error occurred during the lookup
///     (e.g., the library handle is invalid).
/// *   `Err(DlError::InvalidSymbolName)` if the provided symbol name contains an
///     internal null byte, making it invalid for FFI.
#[inline]
pub fn dlsym<S>(handle: &DlHandle, name: S) -> Result<Option<DlSymbol>, DlError>
where
    S: AsRef<OsStr>,
{
    let name_ref = name.as_ref();
    let name_bytes = name_ref.as_bytes();
    let symbol_name = CString::new(name_bytes).map_err(|_| DlError::InvalidSymbolName {
        name: name_ref.to_owned(),
    })?;

    // Clear any existing error state.
    let _ = dlerror();

    // SAFETY: `handle` is guaranteed to be a valid library handle from a `DlHandle`.
    // `symbol_name` is a valid, null-terminated C-string created above.
    // The call is to our securely resolved `dlsym` function.
    let ptr = unsafe { ffi::dlsym(handle.as_raw(), symbol_name.as_ptr()) };

    // Use the safe `ptr::NonNull::new` constructor which returns an Option.
    // This elegantly combines the null-check and the creation of the NonNull pointer.
    match ptr::NonNull::new(ptr) {
        Some(non_null_ptr) => {
            let symbol = DlSymbol {
                _handle: handle.clone(),
                ptr: non_null_ptr,
            };
            Ok(Some(symbol))
        }
        None => {
            if let Some(reason) = dlerror() {
                // A non-null return from `dlerror()` indicates a lookup error.
                let e = DlError::SymbolLookupError {
                    name: name_ref.to_owned(),
                    reason,
                };
                Err(e)
            } else {
                // The symbol was not found.
                Ok(None)
            }
        }
    }
}

/// Translates a memory address to symbolic information.
/// Returns an error if `dladdr` fails.
#[inline]
#[allow(clippy::not_unsafe_ptr_arg_deref)] // `dladdr` is designed to accept any pointer.
pub fn dladdr(addr: *const c_void) -> Result<DlInfo, DlError> {
    let mut dl_info = DlInfo::default();

    // Clear any existing error state.
    let _ = dlerror();

    // SAFETY: `dladdr` safely accepts any pointer for `addr`. We provide a valid mutable pointer to `dl_info`.
    if unsafe { ffi::dladdr(addr, &mut dl_info.0) } == 0 {
        if let Some(reason) = dlerror() {
            // An error message exists, so a real error occurred.
            return Err(DlError::AddressLookupError { addr, reason });
        } else {
            // No error message, so the address was simply not found.
            return Err(DlError::InvalidAddress { addr });
        }
    }

    Ok(dl_info)
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::os::unix::ffi::OsStrExt;

    const LIBM_NAME: &str = "libm.so.6";
    const LIBM_COS_SYMBOL: &str = "cos";
    const LIBC_PRINTF_SYMBOL: &str = "printf";
    const NON_EXISTENT_FILE: &str = "/this/path/absolutely/does/not/exist.so";
    const NON_EXISTENT_SYMBOL: &str = "this_symbol_definitely_does_not_exist";

    #[test]
    fn test_dlopen_success_and_drop() {
        let lib = dlopen(LIBM_NAME, DlOpenFlag::default())
            .expect("Should successfully load a standard library like libm");
        // `lib` is dropped here, and `dlclose` is called automatically.
        drop(lib);
    }

    #[test]
    fn test_dlopen_failure_non_existent_file() {
        let result = dlopen(NON_EXISTENT_FILE, DlOpenFlag::default());
        assert!(
            matches!(result, Err(DlError::LoadFailed { .. })),
            "Expected LoadFailed error, got {:?}",
            result
        );
    }

    #[test]
    fn test_dlopen_failure_invalid_path_with_null_byte() {
        let invalid_path = Path::new("invalid\0path");
        let result = dlopen(invalid_path, DlOpenFlag::default());
        assert!(
            matches!(result, Err(DlError::InvalidPath { .. })),
            "Expected InvalidPath error, got {:?}",
            result
        );
    }

    #[test]
    fn test_dlsym_success_in_loaded_lib() {
        let handle = dlopen(LIBM_NAME, DlOpenFlag::default()).unwrap();
        let cos_symbol = handle
            .find_symbol(LIBM_COS_SYMBOL)
            .expect("Symbol lookup operation itself should succeed")
            .expect("The 'cos' symbol should exist in libm");

        assert!(!cos_symbol.as_raw().is_null());
    }

    #[test]
    fn test_dlsym_failure_non_existent_symbol() {
        let lib = dlopen(LIBM_NAME, DlOpenFlag::default()).unwrap();
        let result = lib.find_symbol(NON_EXISTENT_SYMBOL);
        assert!(
            matches!(result, Ok(None)),
            "Expected Ok(None) for a non-existent symbol, got {:?}",
            result
        );
    }

    #[test]
    fn test_dlsym_failure_invalid_symbol_with_null_byte() {
        let lib = dlopen(LIBM_NAME, DlOpenFlag::default()).unwrap();
        let invalid_symbol = OsStr::from_bytes(b"invalid\0symbol");
        let result = lib.find_symbol(invalid_symbol);
        assert!(
            matches!(result, Err(DlError::InvalidSymbolName { .. })),
            "Expected InvalidSymbolName error, got {:?}",
            result
        );
    }

    #[test]
    fn test_dlsym_with_rtld_next() {
        let printf_symbol = dlsym(&NEXT_HANDLE, LIBC_PRINTF_SYMBOL)
            .expect("Symbol lookup operation should succeed")
            .expect("The 'printf' symbol should be globally available");
        assert!(!printf_symbol.as_raw().is_null());

        let result = dlsym(&NEXT_HANDLE, NON_EXISTENT_SYMBOL)
            .expect("Symbol lookup operation should succeed");
        assert!(
            result.is_none(),
            "Expected None for a non-existent symbol with RTLD_NEXT"
        );
    }

    #[test]
    fn test_dladdr_success_with_known_symbol() {
        let printf_symbol = match dlsym(&NEXT_HANDLE, LIBC_PRINTF_SYMBOL).unwrap() {
            Some(symbol) => symbol,
            None => {
                eprintln!(
                    "Skipping test: could not find '{}' via RTLD_NEXT.",
                    LIBC_PRINTF_SYMBOL
                );
                return;
            }
        };

        let dl_info = dladdr(printf_symbol.as_raw())
            .expect("dladdr call for a known function pointer should succeed");

        let file_name = dl_info
            .file_name()
            .expect("dladdr should find the file name");
        let file_name_str = file_name.to_string_lossy();
        assert!(
            file_name_str.contains("libc") || file_name_str.contains("libSystem"),
            "Expected file name '{}' to be a C or system library",
            file_name_str
        );

        let symbol_name = dl_info
            .symbol_name()
            .expect("dladdr should find the symbol name");
        assert!(symbol_name.to_string_lossy().contains(LIBC_PRINTF_SYMBOL));

        assert!(dl_info.base_addr().is_some());
        assert!(dl_info.symbol_addr().is_some());
    }

    #[test]
    fn test_dladdr_failure_with_dangling_pointer() {
        let result = dladdr(-1isize as *const c_void);
        assert!(
            matches!(result, Err(DlError::InvalidAddress { .. })),
            "Expected InvalidAddress error for a dangling pointer, got {:?}",
            result
        );
    }

    #[test]
    fn test_calling_resolved_function() {
        let lib = dlopen(LIBM_NAME, DlOpenFlag::default()).unwrap();
        let cos_symbol = lib.find_symbol(LIBM_COS_SYMBOL).unwrap().unwrap();

        type CosFn = unsafe extern "C" fn(f64) -> f64;

        let cos_func: CosFn = unsafe { cos_symbol.cast() };
        let result = unsafe { cos_func(0.0) };

        assert!((result - 1.0).abs() < 1e-9, "cos(0.0) should be 1.0");
    }
}
